
import datetime
import os

import pandas as pd    
import baostock as bs   # api 文档： http://baostock.com/baostock/index.php/Python_API%E6%96%87%E6%A1%A3 
import logging
from concurrent.futures import ThreadPoolExecutor, as_completed

"""
说明：A股每日开盘时间4小时。因此，15分钟线数据量只是日线的16倍。
15分钟线可以观察到一天股票的变化趋势，机器学习就可能学习到操盘痕迹。
    注意：15 分钟 K 线：每根 K 线覆盖 15 分钟的交易时间，其 “开盘价” 是这 15 分钟内的第一笔成交价格，“收盘价” 是这 15 分钟内的最后一笔成交价格。

公司的运营能力也是很重要的参考资料：
季频盈利能力，运营能力，成长能力，偿债能力，现金流量，业绩快报，业绩预告等。
由于这个是以季度为单位的，统计时间尺度与日均线不一样。

transform 应该要有整合不同尺度信息的能力。

尺度分4个：
季k线 & 季度盈利数据，   
周k线    
日k线    
15分钟线
"""


'''
日线指标参数包括：'date,code,open,high,low,close,preclose,volume,amount,adjustflag,turn,tradestatus,pctChg,peTTM,pbMRQ,psTTM,pcfNcfTTM,isST'
周、月线指标参数包括：'date,code,open,high,low,close,volume,amount,adjustflag,turn,pctChg'
分钟指标参数包括：'date,time,code,open,high,low,close,volume,amount,adjustflag'

adjustflag：复权类型，默认不复权：3；1：后复权；2：前复权。已支持分钟线、日线、周线、月线前后复权。
'''

# 是否删除停盘数据
# DROP_SUSPENSION = True


class stock_data:

    def __init__(self) -> None:
        logging.basicConfig()
        self.directory_path = './stock_data/'
        #self.default_begin_date = '2025-01-02' # 2010年1月4日，是A股2010年第一个交易日。
        self.default_begin_date = '2010-01-04' # 2010年1月4日，是A股2010年第一个交易日。
        self.default_date = "2025-01-02"   # 2025年1月2日是A股的首个交易日。
        self.default_list_file = self.directory_path + self.default_date+"stock_list.csv"
        self.list_file = self.directory_path + "stock_list.csv"

        self.datas_w = 'date,open,high,low,close,volume,amount,turn,pctChg'
        self.datas_d = 'date,open,high,low,close,volume,amount,turn,pctChg'
        self.datas_15 = 'date,time,open,high,low,close,volume,amount'
        self.datas_map = {
            'w': self.datas_w,
            'd': self.datas_d,
            '15': self.datas_15
        }
        # 判断文件夹路径是否存在，如果不存在则创建路径，并下载默认日期的股票列表。
        if not os.path.exists(self.directory_path):
            os.makedirs(self.directory_path)
            self.list_file = self.default_list_file
            self.update_stock_list(self.default_date)        # 2025年1月2日是A股的首个交易日。
            self.list_file = self.directory_path + "stock_list.csv"

    def update_stock_list(self, date = None):
        """
            更新指定日期的股票列表。  
            只要对比现在的股票列表与十年前的股票列表，就可以知道那些股票存在至今。
        Args:
            date (str, optional): 查询日期，格式YYYY-MM-DD。默认为None表示最新日期。
            另外，非交易日返回的股票列表是空的。
        Returns:
            _list_: A股股票代码列表，每个元素是 [code, code_name]
        """
        
        try:
            lg = bs.login()

            if lg.error_code != '0':
                logging.error(f"baostock登录失败: {lg.error_msg}")
                return []

            # 获取指定日期的指数、股票数据
            stock_rs = bs.query_all_stock(date)
            if stock_rs.error_code != '0':
                logging.error(f"baostock查询股票列表失败: {stock_rs.error_msg}")
                bs.logout()
                return []

            stock_df = stock_rs.get_data() #
            if stock_df.empty:
                logging.warning("baostock返回的股票列表为空")
                bs.logout()
                return []
            
            # 登出BaoStock
            bs.logout()  

            # 下面筛选A股股票
            '''
            A 股市场结构简介
                沪市（Shanghai Stock Exchange）：
                    主板：股票代码以 600、601、603、605 开头（如 600000 浦发银行）。
                    科创板：代码以 688 开头（如 688001 华兴源创），属于沪市特殊板块。
                深市（Shenzhen Stock Exchange）：
                    主板：代码以 000 开头（如 000001 平安银行）。
                    中小板：代码以 002 开头（如 002001 新和成），2021 年后并入深市主板。
                    创业板：代码以 300 开头（如 300001 特锐德）。
            '''
            a_shares = stock_df[
                (stock_df['code'].str.match(r'^sh\.6\d{5}$')) | 
                (stock_df['code'].str.match(r'^sz\.(000|002)\d{3}$'))
                # 正则表达式：^匹配开头，(000|002)分组匹配000或002，\d{3}匹配任意3个数字，$	匹配字符串的结束位置
            ]

            stock_df = a_shares    # a_shares['code'] # 这里是提取股票代码列
            stock_df.to_csv(self.list_file, encoding = 'utf-8', index = False)
            return stock_df[['code','code_name']].values.tolist() 
        except Exception as e:
            logging.exception(f"更新股票列表时发生异常: {str(e)}")
            # 确保在发生异常时登出
            bs.logout()
            return []
        
    def load_stock_list(self):
        """_summary_ 从文件中加载股票列表。           
        Returns:
            list: _A股股票代码列表，每个元素是是 [code, code_name]_
        """
        if not os.path.exists(self.list_file):
            logging.warning("stock_list 文件不存在，返回默认股票列表："+ self.default_list_file)            
            df = pd.read_csv(self.default_list_file)
            return df[['code','code_name']].values.tolist()
        
        df = pd.read_csv(self.list_file)
        return df[['code','code_name']].values.tolist()

    '''
    对于日均线，15min线，数据量大，从零下载的时间很长。
    注意：这里的均线数据是从 self.default_begin_date = '2010-01-04' 开始下载的。从这个点开始计算时间戳。截至日期是当前日期。
    因此，这里要可以保证分段下载：
        1、第一次运行下载了10只股票，第二次运行就从第11只股票开始下载。同时，文件保存为二进制格式，减少内存占用。
        2、隔几天运行的时候，检查每只股票更新的最近日期，从最近日期开始下载文件。
    '''

    def update_stock_data(self,stk_list,frequency = 'd',todate = datetime.date.today()):
        """_summary_ 刷新股票数据，保存为 parquet 二进制文件。这里需要事先登录 baostock 
        Args:
            stk_list (_type_):  股票列表，每个元素是 [code, code_name]
            frequency (str, optional): 均线周期，可选'w'星期，'d'天,'15'15分钟 Defaults to 'd'.
            todate (_type_, optional): 截止日期，默认到今天. Defaults to datetime.date.today().
        """

        # stk_list = self.load_stock_list()[0:4]
        
        # 数据参数配置
        datas = self.datas_map.get(frequency, self.datas_map['w']) # 若键值frequency不存在，用'w'键值

        '''
        周均线 datas 意义；
            date	交易所行情日期	格式：YYYY-MM-DD	
            code	证券代码	格式：sh.600000。sh：上海，sz：深圳	
            open	开盘价格	精度：小数点后4位；单位：人民币元	
            high	最高价	精度：小数点后4位；单位：人民币元	
            low	    最低价	精度：小数点后4位；单位：人民币元	
            close	收盘价	精度：小数点后4位；单位：人民币元	
            volume	成交数量	单位：股	
            amount	成交金额	精度：小数点后4位；单位：人民币元	
            adjustflag	复权状态	不复权、前复权、后复权	
            turn	换手率	精度：小数点后6位；单位：%	
            pctChg	涨跌幅（百分比）	精度：小数点后6位	涨跌幅=[(区间最后交易日收盘价-区间首个交易日前收盘价)/区间首个交易日前收盘价]*100%


        日均线datas的意义：
            date	交易日期（格式：YYYY-MM-DD）
            code	证券代码
            open	开盘价（单位：元）
            high	最高价（单位：元）
            low	最低价（单位：元）
            close	收盘价（单位：元）
            preclose	前收盘价（上一交易日收盘价，单位：元）
            volume	成交量（单位：股）
            amount	成交额（单位：元）
            adjustflag	复权状态（1：前复权；2：后复权；3：不复权）
            turn	换手率（单位：%，计算公式：成交量 / 流通股本）
            tradestatus	交易状态（1：正常交易；0：停牌）
            pctChg	涨跌幅（单位：%，计算公式：(收盘价 - 前收盘价)/ 前收盘价 ×100）
            isST	是否为 ST 股（1：是；0：否）

        分钟均线 datas 的意义： 
        注意：仅支持近 3 年的分钟级数据
            date	交易日期（格式：YYYY-MM-DD）
            time	交易时间（格式：YYYYMMDDHHMMSS，如 20240701093500 表示 9:35）
            code	证券代码
            open	开盘价（单位：元）
            high	最高价（单位：元）
            low	    最低价（单位：元）
            close	收盘价（单位：元）
            volume	成交量（单位：股）
            amount	成交额（单位：元）
            adjustflag	复权类型（1 = 前复权，2 = 后复权，3 = 不复权）     
        '''

        adjustflag = '3'  #复权类型，默认不复权：3；1：后复权；2：前复权。
        '''
        1	前复权	以当前价格为基准，调整历史价格（适合技术分析）
        2	后复权	以历史价格为基准，调整当前价格（适合长期收益计算）
        3	不复权	保持原始价格（适合查看真实成交价格）
        注意：复权会影响 open、high、low、close 等价格字段，volume（成交量）不受复权影响。
        由于数据每日实时更新，因此选择不复权，减少数据刷新次数。
        '''
        # 创建目标目录（可处理多层级目录情况）
        directory_path = self.directory_path + frequency + '/'
        if not os.path.exists(directory_path):
            os.makedirs(directory_path)
        i = 0
        for code_msg in stk_list:
            i=i+1
            logging.info("进度 = " + str(i/len(stk_list)) )
            code = code_msg[0]  # 0是股票代码，1是名字。
            datapath =  directory_path + code + '.parquet'

            if not os.path.exists(datapath): # 文件不存在
                fromdate = self.default_begin_date # 从默认起始日期开始。
                k_rs = bs.query_history_k_data_plus(code, 
                                                    datas, 
                                                    start_date = fromdate, end_date = todate.strftime('%Y-%m-%d'),
                                                    frequency = frequency, adjustflag = adjustflag)
                if k_rs is not None:
                    out_df = k_rs.get_data()
                    # 删除停牌数据
                    if 'volume' in out_df.columns: out_df = out_df[out_df.volume != '0'] 
                    out_df.to_parquet(datapath,engine= "pyarrow")
                    logging.info("下载"+datapath)
                    #  engine= "pyarrow",compression='snappy' 
                    # engine：Parquet 引擎（如 'pyarrow'、'fastparquet'）
                    # compression：压缩方式（如 'snappy'、'gzip'、'none'）,默认即可

            else: # 如果已经有文件了，读取文件最后一列的时间，更新到最近的时间。
                out_df = pd.read_parquet(datapath,engine= "pyarrow")
                fromdate = ''
                if out_df.empty : 
                    fromdate = self.default_begin_date # 从默认起始日期开始。
                    logging.info("无数据，从头下载文件："+datapath)
                else:
                    fromdate = out_df.loc[out_df.index[-1],'date'] # 获取最后一行的日期

                d1 = datetime.datetime.strptime(str(fromdate), '%Y-%m-%d').date()
                d2 = todate
                if( d1+datetime.timedelta(2) < d2 ): # 如果时间不同，才继续执行。
                    k_rs = bs.query_history_k_data_plus(code, 
                                                        datas, 
                                                        start_date = fromdate, end_date = todate.strftime('%Y-%m-%d'),
                                                        frequency = frequency, adjustflag = adjustflag)
                    if k_rs is not None:
                        new_df = k_rs.get_data()
                        # 删除停牌数据
                        if 'volume' in new_df.columns: new_df = new_df[new_df.volume != '0'] 
                        if len(new_df) > 0 : 
                            if out_df.empty : 
                                new_df.to_parquet(datapath,engine= "pyarrow")
                            else:
                                out_df = out_df[out_df.date != fromdate]  # 删除最后一天的数据.
                                out_df = pd.concat([out_df,new_df],ignore_index=True) # ignore_index=true,：忽略原DataFrame的行索引，生成新的连续整数索引（0, 1, 2, ...） 这样两个dataframe 拼接的时候，可以保证行索引唯一。
                                out_df.to_parquet(datapath,engine= "pyarrow")
                            logging.info("更新"+datapath)

                


    def load_stock_data(self,stk_list,frequency = 'd'):
        """_summary_ 从文件中加载均线数据，并放到列表中返回。列表元素是元组：(股票代码, dataframe)
        Args:
            stk_list (_type_):  股票列表，每个元素是 [code, code_name]
            frequency (str, optional): 均线周期，可选'w'星期，'d'天,'15'15分钟 Defaults to 'd'.
        Returns:
            _type_: 列表，元素是元组：(股票代码, dataframe)
        """

        df_list = []
        directory_path = self.directory_path + frequency + '/'
        if not os.path.exists(directory_path):
            logging.warning(directory_path+"路径不存在，请先调用 update_stock_data 加载数据")
            return df_list
        for code_msg in stk_list:
            code = code_msg[0]  # 0是股票代码，1是名字。
            datapath =  directory_path + code + '.parquet'
            if not os.path.exists(datapath):
                logging.warning(datapath+"路径不存在，请先调用 update_stock_data 加载数据")
            else:
                # 加载文件，放到list中。  
                df = pd.read_parquet(datapath,engine= "pyarrow")
                df_list.append( (code, df) )

        return df_list



if __name__ == '__main__':

    logging.basicConfig(level=logging.DEBUG)

    stk_data = stock_data()

    if True: # 测试 update_stock_list，load_stock_list。  -- 下载股票列表，与加载股票列表。
        #a = stk_data.update_stock_list()
        a = stk_data.load_stock_list()
        #print(a)
        ls = stk_data.load_stock_data(a[0:3],'d')
        print(ls)

        for (code ,df) in ls:
            print(df.apply(pd.to_numeric, errors='coerce').to_numpy()[ : ,1 : 6] )
 
    

    if False:  # 测试更新股票均线数据。 单线程
        a = stk_data.load_stock_list()
        bs.login()

        #stk_data.update_stock_data(a[0:3],'d',datetime.date(2025,3,4))

        stk_data.update_stock_data(a[0:3],'d',datetime.date(2025,4,5))
        ls = stk_data.load_stock_data(a[0:3],'d')
        print(ls)

        #stk_data.update_stock_data(a[0:3],'d',datetime.date(2025,5,4))
        bs.logout()

    if False:  # 加载股票数据，到今天的

        a = stk_data.load_stock_list()
        bs.login()

        #stk_data.update_stock_data(a[0:3],'d',datetime.date(2025,3,4))

        #stk_data.update_stock_data(a,'w')
        #stk_data.update_stock_data(a,'d')
        stk_data.update_stock_data(a,'15')

        #ls = stk_data.load_stock_data(a,'d')

        #stk_data.update_stock_data(a[0:3],'d',datetime.date(2025,5,4))
        bs.logout()






